{
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python37664bitsandboxconda88e9e3c6ed2640baa2e35a678edccfd6",
   "display_name": "Python 3.7.6 64-bit ('sandbox': conda)"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import config\n",
    "import torch\n",
    "import numpy as np\n",
    "import torch_model\n",
    "from sklearn.neighbors import NearestNeighbors\n",
    "import torchvision.transforms as T\n",
    "import os\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def load_image_tensor(image_path, device):\n",
    "    image_tensor = T.ToTensor()(Image.open(image_path))\n",
    "    image_tensor = image_tensor.unsqueeze(0)\n",
    "    print(image_tensor.shape)\n",
    "    # input_images = image_tensor.to(device)\n",
    "    return image_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def compute_similar_images(image_path, num_images, embedding, device):\n",
    "    image_tensor = load_image_tensor(image_path, device)\n",
    "    # image_tensor = image_tensor.to(device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        image_embedding = encoder(image_tensor).cpu().detach().numpy()\n",
    "\n",
    "    print(image_embedding.shape)\n",
    "\n",
    "    flattened_embedding = image_embedding.reshape((image_embedding.shape[0], -1))\n",
    "    print(flattened_embedding.shape)\n",
    "\n",
    "    knn = NearestNeighbors(n_neighbors=num_images, metric=\"cosine\")\n",
    "    knn.fit(embedding)\n",
    "\n",
    "    _, indices = knn.kneighbors(flattened_embedding)\n",
    "    indices_list = indices.tolist()\n",
    "    print(indices_list)\n",
    "    return indices_list\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def plot_similar_images(indices_list):\n",
    "    indices = indices_list[0]\n",
    "    for index in indices:\n",
    "        img_name = str(index - 1) + \".jpg\"\n",
    "        img_path = os.path.join(config.DATA_PATH + img_name)\n",
    "        print(img_path)\n",
    "        img = Image.open(img_path).convert(\"RGB\")\n",
    "        plt.imshow(img)\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "TEST_IMAGE_PATH = \"../small_data_oil_for_classification/images/all/4025.png\"\n",
    "NUM_IMAGES = 10\n",
    "ENCODER_MODEL_PATH = \"../data/models/deep_encoder.pt\"\n",
    "EMBEDDING_PATH = \"../data/models/data_embedding_f.npy\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Error(s) in loading state_dict for ConvEncoder:\n\tMissing key(s) in state_dict: \"conv4.weight\", \"conv4.bias\", \"conv5.weight\", \"conv5.bias\". ",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mRuntimeError\u001B[0m                              Traceback (most recent call last)",
      "Input \u001B[1;32mIn [6]\u001B[0m, in \u001B[0;36m<cell line: 5>\u001B[1;34m()\u001B[0m\n\u001B[0;32m      2\u001B[0m encoder \u001B[38;5;241m=\u001B[39m torch_model\u001B[38;5;241m.\u001B[39mConvEncoder()\n\u001B[0;32m      4\u001B[0m \u001B[38;5;66;03m# Load the state dict of encoder\u001B[39;00m\n\u001B[1;32m----> 5\u001B[0m \u001B[43mencoder\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mload_state_dict\u001B[49m\u001B[43m(\u001B[49m\u001B[43mtorch\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mload\u001B[49m\u001B[43m(\u001B[49m\u001B[43mENCODER_MODEL_PATH\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mmap_location\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mdevice\u001B[49m\u001B[43m)\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m      6\u001B[0m encoder\u001B[38;5;241m.\u001B[39meval()\n\u001B[0;32m      7\u001B[0m encoder\u001B[38;5;241m.\u001B[39mto(device)\n",
      "File \u001B[1;32mD:\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py:1497\u001B[0m, in \u001B[0;36mModule.load_state_dict\u001B[1;34m(self, state_dict, strict)\u001B[0m\n\u001B[0;32m   1492\u001B[0m         error_msgs\u001B[38;5;241m.\u001B[39minsert(\n\u001B[0;32m   1493\u001B[0m             \u001B[38;5;241m0\u001B[39m, \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mMissing key(s) in state_dict: \u001B[39m\u001B[38;5;132;01m{}\u001B[39;00m\u001B[38;5;124m. \u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;241m.\u001B[39mformat(\n\u001B[0;32m   1494\u001B[0m                 \u001B[38;5;124m'\u001B[39m\u001B[38;5;124m, \u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;241m.\u001B[39mjoin(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;132;01m{}\u001B[39;00m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;241m.\u001B[39mformat(k) \u001B[38;5;28;01mfor\u001B[39;00m k \u001B[38;5;129;01min\u001B[39;00m missing_keys)))\n\u001B[0;32m   1496\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mlen\u001B[39m(error_msgs) \u001B[38;5;241m>\u001B[39m \u001B[38;5;241m0\u001B[39m:\n\u001B[1;32m-> 1497\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mRuntimeError\u001B[39;00m(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mError(s) in loading state_dict for \u001B[39m\u001B[38;5;132;01m{}\u001B[39;00m\u001B[38;5;124m:\u001B[39m\u001B[38;5;130;01m\\n\u001B[39;00m\u001B[38;5;130;01m\\t\u001B[39;00m\u001B[38;5;132;01m{}\u001B[39;00m\u001B[38;5;124m'\u001B[39m\u001B[38;5;241m.\u001B[39mformat(\n\u001B[0;32m   1498\u001B[0m                        \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m\u001B[38;5;18m__class__\u001B[39m\u001B[38;5;241m.\u001B[39m\u001B[38;5;18m__name__\u001B[39m, \u001B[38;5;124m\"\u001B[39m\u001B[38;5;130;01m\\n\u001B[39;00m\u001B[38;5;130;01m\\t\u001B[39;00m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;241m.\u001B[39mjoin(error_msgs)))\n\u001B[0;32m   1499\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m _IncompatibleKeys(missing_keys, unexpected_keys)\n",
      "\u001B[1;31mRuntimeError\u001B[0m: Error(s) in loading state_dict for ConvEncoder:\n\tMissing key(s) in state_dict: \"conv4.weight\", \"conv4.bias\", \"conv5.weight\", \"conv5.bias\". "
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "encoder = torch_model.ConvEncoder()\n",
    "\n",
    "# Load the state dict of encoder\n",
    "encoder.load_state_dict(torch.load(ENCODER_MODEL_PATH, map_location=device))\n",
    "encoder.eval()\n",
    "encoder.to(device)\n",
    "\n",
    "# Loads the embedding\n",
    "embedding = np.load(EMBEDDING_PATH)\n",
    "\n",
    "indices_list = compute_similar_images(TEST_IMAGE_PATH, NUM_IMAGES, embedding, device)\n",
    "plot_similar_images(indices_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": []
  }
 ]
}